# 本脚本用于测试dpa在线线程调度+线程运行的时间消耗
import subprocess
import os
import csv
import re
import json
from collections import defaultdict

# 配置参数 - 更简洁的可扩展结构
class TestConfig:
    """测试配置类，集中管理所有可配置参数"""
    def __init__(self):
        # 基础配置
        self.results_dir = "../../results_raw/dpa_case_study"
        self.output_file = "dpa_thread_analyze_v5.csv"
        self.test_program = "./../../../bin_dpu/test_operator"
        self.timeout = 120
        
        # 核心可枚举参数 - 易于扩展
        # self.batch_sizes = [8192 * (8 ** i) for i in range(6)]
        self.batch_sizes = [2**i for i in range(0, 25)]   # 1 到 16777216

        # self.batch_sizes = [8192]
        self.thread_nums = [1, 2, 4, 8, 12, 16, 32, 64, 96, 128, 160, 192, 224, 253]
        # self.thread_nums = [128, 253]
        self.dpa_types = [0, 2]
        self.device_types = ["dpa"]  # 支持多种设备类型
        self.op_types = ["sigrid_hash"]  # 可添加更多操作类型
        self.execute_time = [20]  # 添加execute_time参数列表
        
        # 高级配置
        self.save_raw_output = False
        self.raw_output_dir = "../raw_outputs"
    
    def get_total_runs(self):
        """计算总测试次数"""
        return (len(self.batch_sizes) * len(self.thread_nums) * 
                len(self.dpa_types) * len(self.device_types) * 
                len(self.op_types) * len(self.execute_time))  # 添加execute_time

# 初始化配置
CONFIG = TestConfig()

# 时间参数提取器框架
class TimeExtractor:
    def __init__(self, pattern, converter):
        self.pattern = pattern
        self.converter = converter
    
    def extract(self, stdout):
        match = re.search(self.pattern, stdout)
        if not match:
            return None
        
        groups = match.groups()
        try:
            return self.converter(*groups)
        except Exception as e:
            print(f"转换错误: {e}")
            return None

class UnitTimeExtractor(TimeExtractor):
    """处理带单位的时间值（如1.2ms）"""
    def __init__(self, pattern):
        super().__init__(pattern, self.convert_units)
    
    @staticmethod
    def convert_units(value_str, unit):
        value = float(value_str)
        unit = unit.lower()
        if unit == 's': return value * 1_000_000
        if unit == 'ms': return value * 1_000
        if unit == 'us': return value
        if unit == 'ns': return value / 1000
        raise ValueError(f"未知单位: {unit}")

class CycleTimeExtractor(TimeExtractor):
    """处理时钟周期值（如1200 cycles）"""
    def __init__(self, pattern):
        super().__init__(pattern, self.convert_cycles)
    
    @staticmethod
    def convert_cycles(cycles_str):
        cycles = int(cycles_str)
        return cycles  # 保持为原始周期数

# 时间参数配置
class TimeParamManager:
    """管理时间参数提取配置"""
    def __init__(self):
        # 新的时间参数 - 根据您的需求
        self.base_params = {
            "thread_time": UnitTimeExtractor(
                r'max_execute_time_per_thread\s*=\s*([\d.]+)\s*([a-zA-Z]+)'
            ),
            "normal_mem_access": CycleTimeExtractor(
                r'normal_mem_access cycle = (\d+)'
            ),
            "get_host_mem_addr": CycleTimeExtractor(
                r'get_host_mem_addr cycle = (\d+)'
            ),
            "operation_time": CycleTimeExtractor(
                r'operation_time cycle = (\d+)'
            ),
            "atomic_store_time": CycleTimeExtractor(
                r'atomic_store_time cycle = (\d+)'
            ),
            "total_time": CycleTimeExtractor(
                r'total_time cycle = (\d+)'
            )
        }
    
    def get_params(self):
        """获取所有时间参数提取器"""
        return self.base_params

# 创建结果目录
os.makedirs(CONFIG.results_dir, exist_ok=True)
output_path = os.path.join(CONFIG.results_dir, CONFIG.output_file)

# 如果需要保存原始输出，创建输出目录
if CONFIG.save_raw_output:
    os.makedirs(CONFIG.raw_output_dir, exist_ok=True)

# 初始化时间参数管理器
time_param_manager = TimeParamManager()

# 打开文件以写入结果
with open(output_path, 'w', newline='') as csvfile:
    # 固定表头 - 包含所有核心参数
    fieldnames = [
        'op_type', 'device_type', 'batch_size', 'thread_num', 'dpa_type', 'execute_time',
        'thread_time', 'p50_time', 'p99_time', 'thread_ratio',
        'normal_mem_access', 'get_host_mem_addr', 'operation_time', 
        'atomic_store_time', 'total_time'
    ]
    
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    total_runs = CONFIG.get_total_runs()
    current_run = 0
    run_stats = defaultdict(int)  # 用于统计运行情况
    
    # 多层嵌套循环测试所有参数组合
    for op_type in CONFIG.op_types:
        for device_type in CONFIG.device_types:
            for batch_size in CONFIG.batch_sizes:
                for thread_num in CONFIG.thread_nums:
                    for dpa_type in CONFIG.dpa_types:
                        for execute_time in CONFIG.execute_time:  # 添加execute_time循环
                            current_run += 1
                            print(f"运行测试 ({current_run}/{total_runs}): "
                                  f"op_type={op_type}, device={device_type}, "
                                  f"batch_size={batch_size}, thread_num={thread_num}, "
                                  f"dpa_type={dpa_type}, execute_time={execute_time}")
                            
                            # 构建命令（添加execute_time参数）
                            command = [
                                CONFIG.test_program,
                                f'--batch_size={batch_size}',
                                f'--thread_num={thread_num}',
                                f'--device_type={device_type}',
                                f'--op_type={op_type}',
                                f'--dpa_type={dpa_type}',
                                f'--execute_time={execute_time}'  # 添加execute_time参数
                            ]
                            
                            # 执行命令
                            try:
                                result = subprocess.run(
                                    command, 
                                    stdout=subprocess.PIPE, 
                                    stderr=subprocess.PIPE, 
                                    text=True, 
                                    encoding='utf-8', 
                                    errors='ignore', 
                                    timeout=CONFIG.timeout
                                )
                                run_stats['completed'] += 1
                            except subprocess.TimeoutExpired:
                                print(f"超时: op_type={op_type}, device={device_type}, "
                                      f"batch_size={batch_size}, thread_num={thread_num}, "
                                      f"dpa_type={dpa_type}, execute_time={execute_time}")
                                run_stats['timeout'] += 1
                                continue
                            except Exception as e:
                                print(f"执行错误: {e}")
                                run_stats['error'] += 1
                                continue

                            stdout = result.stdout
                            stderr = result.stderr
                            
                            # 保存原始输出用于调试
                            if CONFIG.save_raw_output:
                                output_filename = (
                                    f"output_{op_type}_{device_type}_bs{batch_size}_"
                                    f"th{thread_num}_dpa{dpa_type}_exe{execute_time}.txt"
                                )
                                output_path = os.path.join(CONFIG.raw_output_dir, output_filename)
                                with open(output_path, 'w') as f:
                                    f.write(f"Command: {' '.join(command)}\n\n")
                                    f.write("="*50 + " STDOUT " + "="*50 + "\n")
                                    f.write(stdout)
                                    f.write("\n\n" + "="*50 + " STDERR " + "="*50 + "\n")
                                    f.write(stderr)
                            
                            if stderr:
                                print(f"错误输出: {stderr}")
                                run_stats['stderr'] += 1
                            
                            # 1. 提取所有时间参数
                            time_values = {}
                            missing_params = []
                            
                            # 获取参数提取器
                            extractors = time_param_manager.get_params()
                            
                            for param_name, extractor in extractors.items():
                                value = extractor.extract(stdout)
                                if value is not None:
                                    time_values[param_name] = value
                                else:
                                    missing_params.append(param_name)
                            
                            # 检查必要参数
                            if missing_params:
                                print(f"警告: 缺少参数 {', '.join(missing_params)} - "
                                      f"op_type={op_type}, device={device_type}, "
                                      f"batch_size={batch_size}, thread_num={thread_num}, "
                                      f"dpa_type={dpa_type}, execute_time={execute_time}")
                                run_stats['missing_params'] += 1
                                continue
                            
                            # 2. 提取时间统计表格
                            detected_op_type = None
                            p50_time = None
                            p99_time = None
                            
                            # 使用正则表达式匹配表格行
                            table_pattern = r'\| (\w+)\s+\| (\d+\.?\d*)\s+(\w+)\s+\| (\d+\.?\d*)\s+(\w+)\s+\|'
                            for line in stdout.splitlines():
                                match = re.search(table_pattern, line)
                                if match:
                                    detected_op_type = match.group(1)
                                    p50_value = float(match.group(2))
                                    p50_unit = match.group(3)
                                    p99_value = float(match.group(4))
                                    p99_unit = match.group(5)
                                    
                                    # 转换为微秒
                                    unit_extractor = UnitTimeExtractor(None)
                                    p50_time = unit_extractor.convert_units(p50_value, p50_unit)
                                    p99_time = unit_extractor.convert_units(p99_value, p99_unit)
                                    break
                            
                            if detected_op_type is None or p50_time is None or p99_time is None:
                                print(f"警告: 未找到时间统计表格 - "
                                      f"op_type={op_type}, device={device_type}, "
                                      f"batch_size={batch_size}, thread_num={thread_num}, "
                                      f"dpa_type={dpa_type}, execute_time={execute_time}")
                                run_stats['missing_table'] += 1
                                continue
                            
                            # 验证检测到的操作类型是否匹配
                            if detected_op_type != op_type:
                                print(f"警告: 检测到的操作类型不匹配 "
                                      f"({detected_op_type} != {op_type}) - "
                                      f"device={device_type}, batch_size={batch_size}, "
                                      f"thread_num={thread_num}, dpa_type={dpa_type}, execute_time={execute_time}")
                            
                            # 计算thread_ratio
                            thread_ratio = int(round((time_values["thread_time"] / p50_time) * 100, 0))
                            thread_ratio = max(0, min(100, thread_ratio))
                            
                            # 构建CSV行数据（添加execute_time字段）
                            row_data = {
                                'op_type': op_type,
                                'device_type': device_type,
                                'batch_size': batch_size,
                                'thread_num': thread_num,
                                'dpa_type': dpa_type,
                                'execute_time': execute_time,  # 添加execute_time
                                'thread_time': time_values["thread_time"],
                                'p50_time': p50_time,
                                'p99_time': p99_time,
                                'thread_ratio': thread_ratio,
                                'normal_mem_access': time_values["normal_mem_access"],
                                'get_host_mem_addr': time_values["get_host_mem_addr"],
                                'operation_time': time_values["operation_time"],
                                'atomic_store_time': time_values["atomic_store_time"],
                                'total_time': time_values["total_time"]
                            }
                            
                            # 写入CSV行
                            writer.writerow(row_data)
                            run_stats['success'] += 1

# 打印运行统计信息
print("\n测试完成! 结果已保存到", output_path)
print("运行统计:")
print(f"  总测试次数: {total_runs}")
print(f"  成功完成: {run_stats['success']}")
print(f"  超时: {run_stats.get('timeout', 0)}")
print(f"  错误: {run_stats.get('error', 0)}")
print(f"  标准错误输出: {run_stats.get('stderr', 0)}")
print(f"  缺少参数: {run_stats.get('missing_params', 0)}")
print(f"  缺少时间表格: {run_stats.get('missing_table', 0)}")

# 保存配置信息（添加execute_time）
config_summary = {
    "batch_sizes": CONFIG.batch_sizes,
    "thread_nums": CONFIG.thread_nums,
    "dpa_types": CONFIG.dpa_types,
    "device_types": CONFIG.device_types,
    "op_types": CONFIG.op_types,
    "execute_time": CONFIG.execute_time,  # 添加execute_time
    "run_stats": dict(run_stats)
}

config_path = os.path.join(CONFIG.results_dir, "test_config_summary.json")
with open(config_path, 'w') as f:
    json.dump(config_summary, f, indent=2)

print(f"配置摘要已保存到: {config_path}")
